#!/usr/bin/env python
import math
import sys
import  scipy.integrate as sciint
import numpy
from mpmath import *
import mpmath
import cmath

from unit import *
from bath import *
from constant import *
from propagator import *
from spectral_density import *
from tensor import *
from matrix_index_2d import *
from global_parameter import *
from system_2d import *




def generate_matrix_T_for_kmax1(T, J, stepsize, w0, env_parameters):
	M = [None for temp in range(0, 4)]
	for i in range(0, len(M)):
		M[i] = [None for temp in range(0, 4)]
	for i in range(0, len(M)):
		for j in range(0, len(M[0])):
			M[i][j] = propagator(generate_vector(i)[0], generate_vector(i)[1], generate_vector(j)[0], generate_vector(j)[1],w0, stepsize) * Feyman_Vernon_Specific_Distance_Contribution(T, J, stepsize, env_parameters, 0, generate_two_indice(j)[0], generate_two_indice(j)[1], generate_two_indice(j)[0], generate_two_indice(j)[1]) * Feyman_Vernon_Specific_Distance_Contribution(T, J, stepsize, env_parameters, 1, generate_two_indice(j)[0], generate_two_indice(j)[1], generate_two_indice(i)[0], generate_two_indice(i)[1])
	return matrix(M)


def tensor_calculate_for_kmax1(system_parameters, inital_state, temperature, Spectral_Density, step, env_parameters, cal_range, recal_w0, pulse_signal = 0, pulse_time = -1, delta0_extra = 0, epsi_extra = 0):
	results = []
	A = inital_state
	# generate detailed system parameter
	delta0 = system_parameters[0]
	epsi = system_parameters[1]
	delta0_deviation = system_parameters[2]
	epsi_deviation = system_parameters[3]
	w0 = system_matrix(delta0, epsi, delta0_deviation, epsi_deviation) 
	T = generate_matrix_T_for_kmax1(temperature, Spectral_Density, step,  w0, env_parameters)
	up_bound = int(cal_range[1] / step + 1)
	N = 0
	pulse_flag = None
	results.append(["time", "spin_up_prob", "spin_down_prob", "thermal_average", "decoherence", "ratio"])
	results.append([1.0 * N * step, re(A[0]), re(A[1]), re(A[0]) - re(A[1]), fabs(A[2]), 0])
	while N < up_bound: 
		N += 1
		A = T * A
		result = matrix([0, 0, 0, 0])	
		for i in range(0, len(result)):
			result[i] = A[i] * Feyman_Vernon_Specific_Distance_Contribution(temperature, Spectral_Density, step, env_parameters, 0, generate_two_indice(i)[0], generate_two_indice(i)[1], generate_two_indice(i)[0], generate_two_indice(i)[1])
		results.append([1.0 * N * step, re(result[0]), re(result[1]), re(result[0]) - re(result[1]), fabs(result[2]), re(result[0]) - results[-1][1]])
		if recal_w0 == 1:
			if pulse_signal == 1:
				# Currently, it means that pulse_signal effect cannont be combined with static disorder
				print "Internal Error\n"
				exit(-1)	
			w0 = system_matrix(delta0, epsi, delta0_deviation, epsi_deviation) 
			T = generate_matrix_T_for_kmax1(temperature, Spectral_Density, step,  w0, env_parameters)
		else:
			if pulse_signal == 1:
				if pulse_flag == None:
					w_extra = system_extra_matrix(delta0_extra, epsi_extra, N * step, step, pulse_time)
					if w_extra != 0: 
						T = generate_matrix_T_for_kmax1(temperature, Spectral_Density, step, w0 + w_extra, env_parameters)	
						pulse_flag = True
				elif pulse_flag == True:
					T = generate_matrix_T_for_kmax1(temperature, Spectral_Density, step,  w0, env_parameters)
					pulse_flag = False
				else:
					pass	
	return results

